/*
 * Copyright(C) 2022. Huawei Technologies Co.,Ltd. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */


#include <fstream>
#include <iostream>
#include "BertClassification.h"

namespace {
    const uint32_t MAX_LENGTH = 300;
    const float TEXT_START_CODE = 101.0;
    const float TEXT_END_CODE = 102.0;
    const float TEXT_NOT_FOUND_CODE = 100.0;
    const int FLOAT32_BYTES = 4;
    const uint32_t LABEL_NUMBER = 5;
}

APP_ERROR InitTokenMap(const std::string &vocabTextPath, std::map<std::string, int> &tokenMap)
{
    const std::string text;
    std::ifstream infile;
    // Open label file.
    infile.open(vocabTextPath, std::ios_base::in);
    if(!infile.is_open ()) {
        std::cout << "Open " << vocabTextPath << " file failure!" << std::endl;
        return APP_ERR_COMM_OPEN_FAIL;
    }

    std::string s;
    int count = 0;
    while (std::getline(infile, s)) {
        tokenMap.insert(std::pair<std::string, int>(s, count));
        count++;
    }
    infile.close();
    return APP_ERR_OK;
}

std::vector<std::string> SplitCinese(std::string s)
{
    std::vector<std::string> words;
    for (size_t i = 0; i < s.length();) {
        int cplen = 1;
        if ((s[i] & 0xf8) == 0xf0)
            cplen = 4;
        else if ((s[i] & 0xf0) == 0xe0)
            cplen = 3;
        else if ((s[i] & 0xe0) == 0xc0)
            cplen = 2;
        if ((i + cplen) > s.length())
            cplen = 1;
        words.push_back(s.substr(i, cplen));
        i += cplen;
    }
    return words;
}

APP_ERROR LoadLabels(const std::string &labelPath, std::map<int, std::string> &labelMap)
{
    std::ifstream infile;
    // Open label file.
    infile.open(labelPath, std::ios_base::in);
    std::string s;
    // Check label file validity.
    if (infile.fail()) {
        LogError << "Failed to open label file: " << labelPath << ".";
        return APP_ERR_COMM_OPEN_FAIL;
    }
    labelMap.clear();
    // Construct label map.
    int count = 0;
    while (std::getline(infile, s)) {
        size_t eraseIndex = s.find_last_not_of("\r\n\t");
        if (eraseIndex != std::string::npos) {
            s.erase(eraseIndex + 1, s.size() - eraseIndex);
        }
        labelMap.insert(std::pair<int, std::string>(count, s));
        count++;
    }
    infile.close();
    return APP_ERR_OK;
}

APP_ERROR BertClassification::TextToTensor(const std::string &text, const InitParam &initParam,
                                           std::vector<MxBase::Tensor> &inputs)
{
    APP_ERROR ret;
    std::vector<uint32_t> shape = {1, MAX_LENGTH};
    std::map<std::string, int>::iterator iter;
    uint32_t i, value, size, end_index;
    std::vector<std::string> words = SplitCinese(text);

    float* tensor1Data = new float [MAX_LENGTH];
    // Init data.
    for (i = 0;i < MAX_LENGTH; i++) {
        tensor1Data[i] = 0.0;
    }

    size = words.size();
    if (size > MAX_LENGTH) {
        // Remove start and end characters, length is 2.
        size = MAX_LENGTH - 2;
        end_index = MAX_LENGTH - 1;
    } else {
        end_index = size - 1;
    }

    std::map<std::string, int> tokenMap = {};
    ret = InitTokenMap(initParam.vocabTextPath, tokenMap);
    if (ret != APP_ERR_OK) {
        LogError << "Failed to load tokenMap, ret=" << ret << ".";
        return ret;
    }
    tensor1Data[0] = TEXT_START_CODE;
    // Text decode.
    for (i = 0;i < size; i++) {
        iter = tokenMap.find(words[i]);
        if (iter != tokenMap.end()) {
            value = iter->second;
            tensor1Data[i+1] = float(value);
        } else {
            tensor1Data[i+1] = TEXT_NOT_FOUND_CODE;
        }
    }
    tensor1Data[end_index] = TEXT_END_CODE;

    MxBase::Tensor tensor1((void*)tensor1Data, shape,MxBase::TensorDType::FLOAT32, initParam.deviceId);
    inputs.push_back(tensor1);

    // Init second tensor.
    float* tensor2Data = new float [MAX_LENGTH];
    for (i = 0;i < MAX_LENGTH; i++) {
        tensor2Data[i] = 0.0;
    }
    MxBase::Tensor tensor2((void*)tensor2Data, shape,MxBase::TensorDType::FLOAT32, initParam.deviceId);
    inputs.push_back(tensor2);
    // Release memory.
    delete[] tensor1Data;
    delete[] tensor2Data;
    return APP_ERR_OK;
}

APP_ERROR BertClassification::PostProcess(const std::vector<MxBase::Tensor> &outputs, const InitParam &initParam,
                                          std::string &label)
{
    APP_ERROR ret;
    uint32_t i;
    uint32_t maxIndex = 0;
    float maxValue = 0;
    MxBase::Tensor outTensor = outputs[0];
    outTensor.ToHost();
    // Inference result is tensor(1*5).
    float* inferResult = (float *)outTensor.GetData();
    // Find the category with the highest probability.
    for (i = 0; i < LABEL_NUMBER; i++) {
        if (inferResult[i] > maxValue) {
            maxValue = inferResult[i];
            maxIndex = i;
        }
    }
    // Get label.
    std::map<int, std::string> labelMap = {};
    ret = LoadLabels(initParam.labelPath,labelMap);
    if (ret != APP_ERR_OK) {
        LogError << "Failed to load labels, ret=" << ret << ".";
        return ret;
    }
    label = labelMap.at(maxIndex);
    return APP_ERR_OK;
}



